# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Full connected neighbour list
"""

from typing import Tuple
import torch
from torch import Tensor
from torch.nn import Module

from ..function import reduce_all, get_tensor


class FullConnectNeighbours(Module):
    r"""Full connected neighbour list

    Args:

        num_atoms (int):    Number of atoms

    Supported Platforms:

        ``Ascend`` ``GPU``

    """

    def __init__(self, num_atoms: int):
        super().__init__()
        self.num_atoms = num_atoms
        self.num_neighbours = num_atoms - 1

        self.fc_idx = None
        self.no_idx = None
        self.idx_mask = None
        self.shape = ()
        self.fc_mask = None

        if self.num_neighbours > 0:
            # neighbours for no connection (A*N)
            # (A, 1)
            no_idx = get_tensor(torch.arange(self.num_atoms).reshape(-1, 1))

            # (N)
            nrange = get_tensor(torch.arange(self.num_neighbours))

            # neighbours for full connection (A, N)
            # [[1, 2, 3, ..., N],
            #  [0, 2, 3, ..., N],
            #  [0, 1, 3, ..., N],
            #  ...,
            #  [0, 1, 2, ..., N-1]]
            fc_idx = nrange + (no_idx <= nrange).to(torch.int32)
            no_idx = torch.broadcast_to(
                no_idx, (self.num_atoms, self.num_neighbours))
            idx_mask = fc_idx > no_idx

            # (1, A, N)
            self.fc_idx = fc_idx.unsqueeze(0)
            self.no_idx = no_idx.unsqueeze(0)
            self.idx_mask = idx_mask.unsqueeze(0)

            self.shape = (self.num_atoms, self.num_neighbours)
            self.fc_mask = torch.broadcast_to(get_tensor(True), (1,)+self.shape)

    def set_exclude_index(self, exclude_index: Tensor) -> Tensor:
        """Dummy"""
        return exclude_index

    def print_info(self):
        """print information"""
        info = {'Number of Neighbours': self.num_neighbours}
        return info

    def forward(self,
                  atom_mask: Tensor = None,
                  exclude_index: Tensor = None
                  ) -> Tuple[Tensor, Tensor]:
        r"""Calculate the full connected neighbour list.

        Args:
            atom_mask (Tensor):     Tensor of shape (B, A). Data type is bool.
            exclude_index (Tensor): Tensor of shape (B, A, Ex). Data type is int.

        Returns:
            neighbours (Tensor):    Tensor of shape (B, A, N). Data type is int.
            neighbour_mask (Tensor) Tensor of shape (B, A, N). Data type is bool.

        Symbols:
            B:  Batch size.
            A:  Number of atoms in system.
            N:  Number of neighbour atoms.
            D:  Dimension of position coordinates.
            Ex: Maximum number of excluded neighbour atoms.

        """

        if self.fc_idx is None:
            return None, None

        if atom_mask is None:
            neighbours = self.fc_idx
            neighbour_mask = self.fc_mask
            no_idx = self.no_idx
        else:
            # (B, 1, N)
            mask0 = atom_mask[:, :-1].unsqueeze(-2)
            mask1 = atom_mask[:, 1:].unsqueeze(-2)

            # (B, A, N)
            neighbour_mask = torch.where(self.idx_mask, mask1, mask0)
            neighbour_mask = torch.logical_and(atom_mask.unsqueeze(-1), neighbour_mask)
            fc_idx = torch.broadcast_to(self.fc_idx, neighbour_mask.shape)
            no_idx = torch.broadcast_to(self.no_idx, neighbour_mask.shape)
            neighbours = torch.where(neighbour_mask, fc_idx, no_idx)

        if exclude_index is not None:
            # (B, A, N, Ex) <- (B, A, N, 1) vs (B, A, 1, Ex)
            exc_mask = neighbours.unsqueeze(-1) != exclude_index.unsqueeze(-2)
            # (B, A, N)
            exc_mask = reduce_all(exc_mask, -1)
            neighbour_mask = torch.logical_and(neighbour_mask, exc_mask)
            if atom_mask is None:
                no_idx = torch.broadcast_to(no_idx, neighbour_mask.shape)
            neighbours = torch.where(neighbour_mask, neighbours, no_idx)

        return neighbours, neighbour_mask
